import os
from random import shuffle

import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.signal
from PIL import Image
from torch.utils.data.dataset import Dataset


def letterbox_image(image, label, size):
    label = Image.fromarray(np.array(label))
    '''resize image with unchanged aspect ratio using padding'''
    iw, ih = image.size
    w, h = size
    scale = min(w / iw, h / ih)  # 使一边一样，另一边比目标小
    nw = int(iw * scale)  # 放大后大小
    nh = int(ih * scale)
    # 进行图片的放大#############################这里应该会出问题，我们那里对应的是png为原图，bmp为标记
    image = image.resize((nw, nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128, 128, 128))
    new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))  # 添加padding

    label = label.resize((nw, nh), Image.NEAREST)
    new_label = Image.new('L', size, (0))
    new_label.paste(label, ((w - nw) // 2, (h - nh) // 2))

    return new_image, new_label


def rand(a=0, b=1):
    return np.random.rand() * (b - a) + a


class DeeplabDataset(Dataset):  # 这个要大范围改动
    def __init__(self, train_lines, image_size, num_classes, random_data, dataset_path):
        super(DeeplabDataset, self).__init__()

        self.train_lines = train_lines
        self.train_batches = len(train_lines)  # 说是batch但是感觉用num更好
        self.image_size = image_size  # 想要输出的图片的大小
        self.num_classes = num_classes
        self.random_data = random_data  # Bool
        self.dataset_path = dataset_path

    def __len__(self):
        return self.train_batches

    def rand(self, a=0, b=1):  # 返回[a,b)范围的随机数
        return np.random.rand() * (b - a) + a

    def get_random_data(self, image, label, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5):
        label = Image.fromarray(np.array(label))

        h, w = input_shape
        # resize image
        rand_jit1 = rand(1 - jitter, 1 + jitter)
        rand_jit2 = rand(1 - jitter, 1 + jitter)
        new_ar = w / h * rand_jit1 / rand_jit2

        scale = rand(0.5, 1.5)
        if new_ar < 1:
            nh = int(scale * h)
            nw = int(nh * new_ar)
        else:
            nw = int(scale * w)
            nh = int(nw / new_ar)
        image = image.resize((nw, nh), Image.BICUBIC)
        label = label.resize((nw, nh), Image.NEAREST)
        label = label.convert("L")

        # flip image or not
        flip = rand() < .5
        if flip:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
            label = label.transpose(Image.FLIP_LEFT_RIGHT)

        # place image
        dx = int(rand(0, w - nw))
        dy = int(rand(0, h - nh))
        new_image = Image.new('RGB', (w, h),(128, 128, 128))#这里改动了
        new_label = Image.new('L', (w, h), (0))
        new_image.paste(image, (dx, dy))
        new_label.paste(label, (dx, dy))
        image = new_image
        label = new_label

        # distort image
        hue = rand(-hue, hue)
        sat = rand(1, sat) if rand() < .5 else 1 / rand(1, sat)
        val = rand(1, val) if rand() < .5 else 1 / rand(1, val)
        x = cv2.cvtColor(np.array(image, np.float32) / 255, cv2.COLOR_RGB2HSV)
        x[..., 0] += hue * 360
        x[..., 0][x[..., 0] > 1] -= 1
        x[..., 0][x[..., 0] < 0] += 1
        x[..., 1] *= sat
        x[..., 2] *= val
        x[x[:, :, 0] > 360, 0] = 360
        x[:, :, 1:][x[:, :, 1:] > 1] = 1
        x[x < 0] = 0
        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB) * 255
        return image_data, label

    def __getitem__(self, index):  # 返回单个对象（包含图片和标记）
        if index == 0:
            shuffle(self.train_lines)  # 暂时不清楚在干什么
        # # 获取图片名
        # annotation_line = self.train_lines[index]
        # name = annotation_line.split()[0]

        # # 从文件中读取图像
        # jpg = Image.open(os.path.join(os.path.join(self.dataset_path, "JPEGImages"), name + ".jpg"))  # 图片
        # png = Image.open(os.path.join(os.path.join(self.dataset_path, "SegmentationClass"), name + ".png"))  # 标记
        ###############################这里要根据实际情况进行更改，下一个数据集要9个小时左右，而且还损坏了，所以我是靠感觉写的
        jpg=Image.open(self.train_lines[index][0])
        png=Image.open(self.train_lines[index][1])
        if self.random_data:  # 获取随机变化的图片
            jpg, png = self.get_random_data(jpg, png, (int(self.image_size[1]), int(self.image_size[0])))
        else:  # 固定变化的图片
            jpg, png = letterbox_image(jpg, png, (int(self.image_size[1]), int(self.image_size[0])))

        png = np.array(png)  # 标记转化为np
        png[png >= self.num_classes] = self.num_classes  # 不在类里面的标记转化为最后一个类（最后一个类一般是背景）

        # -------------------------------------------------------#
        #   转化成one_hot的形式
        #   在这里需要+1是因为voc数据集有些标签具有白边部分
        #   我们需要将白边部分进行忽略，+1的目的是方便忽略。
        #   到时候再看一下数据集有无白边决定要不要+1
        # -------------------------------------------------------#
        seg_labels = np.eye(self.num_classes + 1)[png.reshape([-1])]  # 将标记展平，每个代表类的数字用onehot编码代替
        seg_labels = seg_labels.reshape(
            (int(self.image_size[1]), int(self.image_size[0]), self.num_classes + 1))  # 变回原来的形状

        jpg = np.transpose(np.array(jpg), [2, 0, 1]) / 255  # 变回tensor的通道排列顺序并进行归一化

        return jpg, png, seg_labels


# DataLoader中collate_fn使用
def deeplab_dataset_collate(batch):  # 将每个的三部分分别合并
    images = []
    pngs = []
    seg_labels = []
    for img, png, labels in batch:
        images.append(img)
        pngs.append(png)
        seg_labels.append(labels)
    images = np.array(images)
    pngs = np.array(pngs)
    seg_labels = np.array(seg_labels)
    return images, pngs, seg_labels


class LossHistory():  # 存loss的信息和画loss的变化图
    def __init__(self, log_dir):
        import datetime
        curr_time = datetime.datetime.now()
        time_str = datetime.datetime.strftime(curr_time, '%Y_%m_%d_%H_%M_%S')
        self.log_dir = log_dir
        self.time_str = time_str
        self.save_path = os.path.join(self.log_dir, "loss_" + str(self.time_str))
        self.losses = []
        self.val_loss = []

        os.makedirs(self.save_path)

    def append_loss(self, loss, val_loss):
        self.losses.append(loss)
        self.val_loss.append(val_loss)
        with open(os.path.join(self.save_path, "epoch_loss_" + str(self.time_str) + ".txt"), 'a') as f:
            f.write(str(loss))
            f.write("\n")
        with open(os.path.join(self.save_path, "epoch_val_loss_" + str(self.time_str) + ".txt"), 'a') as f:
            f.write(str(val_loss))
            f.write("\n")
        self.loss_plot()

    def loss_plot(self):
        iters = range(len(self.losses))

        plt.figure()
        plt.plot(iters, self.losses, 'red', linewidth=2, label='train loss')
        plt.plot(iters, self.val_loss, 'coral', linewidth=2, label='val loss')
        try:
            if len(self.losses) < 25:
                num = 5
            else:
                num = 15

            plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle='--', linewidth=2,
                     label='smooth train loss')
            plt.plot(iters, scipy.signal.savgol_filter(self.val_loss, num, 3), '#8B4513', linestyle='--', linewidth=2,
                     label='smooth val loss')
        except:
            pass

        plt.grid(True)
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend(loc="upper right")

        plt.savefig(os.path.join(self.save_path, "epoch_loss_" + str(self.time_str) + ".png"))
